import xml.etree.ElementTree as ET

types = ['正常', '扎洞', '毛斑', '擦洞', '毛洞', '织稀', '吊经', '缺经', '跳花', '油／污渍',
         '其他']


def type2id(flaw_type, bin_classify=False):
    """
    传入瑕疵类型名称，返回瑕疵类型编号
    | 0  | 1  | 2  | 3  | 4 | 5  | 6 | 7  | 8 | 9     |10 |
    |正常|扎洞|毛斑|擦洞|毛洞|织稀|吊经|缺经|跳花|油／污渍|其他|

    如果bin_classify==True
    | 0  | 1 |
    |正常|瑕疵|
    :param flaw_type: 字符串类型
    :param bin_classify: bool,是否是二分类问题
    :return: 整型
    """
    if bin_classify:
        return 0 if flaw_type == '正常' else 1

    if flaw_type == "油渍" or flaw_type == "污渍":
        flaw_type = '油／污渍'
    try:
        return types.index(flaw_type)
    except ValueError:
        return types.index('其他')


def typeid2defect_code(typeid):
    """
    传入瑕疵编号，返回defect_code，为最后结果统计服务
    0 --> norm
    1 --> defect_1
    2 --> defect_2
    ...
    :param typeid:  整型
    :return: 字符串
    """
    return 'norm' if typeid == 0 else 'defect_{}'.format(typeid)


class Flaw(object):
    """瑕疵类，仅包含3个属性
        pic_name : 该瑕疵所属的图片名称
        id : 瑕疵类型id
        type : 瑕疵类型名称
        poses: 瑕疵位置，为一个字典数组，一个样本可能有多个瑕疵{xmin: ,ymin: ,xmax: ,ymax: }
    """

    def __init__(self, pic_name=None):
        """
        初始化一个瑕疵类型，如果填写了图片名称，则自动解析相关xml文件，获取瑕疵信息
        :param pic_name: 图片名称
        """
        self.pic_name = pic_name
        self.id = 0
        self.type = "正常"
        self.poses = None

        if pic_name:
            # 通过pic_name获取xml_name
            xml_name = pic_name[:-3] + "xml"
            try:
                # 解析瑕疵文件
                tree = ET.parse(xml_name)
                root = tree.getroot()
                self.pic_name = root[0].text
                self.type = root[4][0].text
                self.id = type2id(self.type)
                self.poses = []
                # 解析所有瑕疵位置
                for flow_pos in root.iter('bndbox'):
                    d = {'xmin': int(flow_pos[0].text),
                         'ymin': int(flow_pos[1].text),
                         'xmax': int(flow_pos[2].text),
                         'ymax': int(flow_pos[3].text)}
                    self.poses.append(d)
            except FileNotFoundError:
                pass

    def get_size(self):
        """
        获取每个瑕疵对应大小
        :return: Yield
        """
        if self.poses:
            for pos in self.poses:
                yield (pos['xmax'] - pos['xmin']) * (pos['ymax'] - pos['ymin'])

    def print_flaw(self):
        print("所属图片名：", self.pic_name)
        print("瑕疵类型：", self.type)
        print("瑕疵类型编号：", self.id)
        print("瑕疵位置：", self.poses)
        print("瑕疵大小：", list(self.get_size()))
